import os
import pickle
import numpy as np
from tensorflow.keras.applications import VGG16
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.models import Model


class ImageSearchEngine:
    def __init__(self, model_path='models'):
        """初始化搜索引擎"""
        # 加载特征提取模型
        base_model = VGG16(weights='imagenet')
        self.feature_model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output)

        # 加载最近邻模型和图片路径
        with open(os.path.join(model_path, 'neighbors.pkl'), 'rb') as f:
            self.neighbors = pickle.load(f)

        with open(os.path.join(model_path, 'image_paths.pkl'), 'rb') as f:
            self.image_paths = pickle.load(f)

    def extract_features(self, img_path):
        """从单张图片提取特征"""
        img = image.load_img(img_path, target_size=(224, 224))
        x = image.img_to_array(img)
        x = np.expand_dims(x, axis=0)
        x = preprocess_input(x)
        features = self.feature_model.predict(x)
        return features.flatten()

    def search(self, query_img_path, top_k=5):
        """搜索相似图片"""
        # 提取查询图片特征
        query_feat = self.extract_features(query_img_path)

        # 搜索最近邻
        distances, indices = self.neighbors.kneighbors([query_feat])

        # 返回结果
        results = []
        for i in range(top_k):
            if i < len(indices[0]):
                img_path = self.image_paths[indices[0][i]]
                distance = distances[0][i]
                similarity = 1 - distance  # 计算相似度（距离越小越相似）
                results.append({
                    'path': img_path,
                    'distance': float(distance),
                    'similarity': float(similarity)  # 添加相似度
                })

        return results

if __name__ == '__main__':
    # 初始化搜索引擎
    search_engine = ImageSearchEngine(model_path='models')  # 指定模型路径

    # 搜索相似图片
    query_path = os.path.join(os.path.dirname(__file__), 'query')
    print(query_path)
    results = search_engine.search(
        query_img_path=  os.path.join(query_path,"FZ1517-500.jpg"),
        top_k=1  # 返回前5个相似结果
    )

    # 打印结果
    for idx, result in enumerate(results):
        print(result)
        print(f"结果 {idx + 1}:")
        print(f"图片路径: {result['path']}")
        print(f"相似度: {result['similarity']:.2f}")
        print("-" * 50)